package com.se.controller;

import com.se.exception.AIRelatedException;
import com.se.factory.ChatModelFactory;
import com.se.service.ConversationService;
import jakarta.servlet.http.HttpSession;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;
import org.springframework.ai.chat.messages.Message;
import java.nio.charset.StandardCharsets;
import java.nio.charset.Charset;

import java.util.Collections;
import java.util.List;
import java.util.Map;

@RestController
@RequestMapping("/ai")
public class AIController {
    private final ConversationService conversationService;
    private final ChatModelFactory chatModelFactory;
    private static final Logger logger = LoggerFactory.getLogger(AIController.class);

    public record StreamChunk(
            String text,
            boolean isComplete,
            String finishReason
    ) {}

    @Autowired
    public AIController(ChatModelFactory chatModelFactory, ConversationService conversationService) {
        this.chatModelFactory = chatModelFactory;
        this.conversationService = conversationService;
    }

    @PostMapping("/generate")
    public Map<String, String> generate(@RequestBody Map<String, String> request, HttpSession session) {

        String apiId = request.get("apiId");
        String message = request.get("message");
        String apiKey = request.get("apiKey");
        String baseUrl = request.get("baseUrl");
        String model = request.get("model");
        double temperature = Double.parseDouble(request.get("temperature"));
        ChatModel chatModel = chatModelFactory.getOrCreate(apiId, apiKey, baseUrl, model, temperature);

        String sessionID = session.getId();
        conversationService.addUserMessage(sessionID, message);
        List<Message> history;
        if(request.containsKey("primary_prompt")) {
            history = conversationService.getOrCreateConversation(sessionID, request.get("primary_prompt"));
        }
        else {
            history = conversationService.getOrCreateConversation(sessionID);
        }
        String response;
        try {
            response = chatModel.call(new Prompt(history)).getResult().getOutput().getText();
            conversationService.addAIResponse(sessionID, response);
        } catch (Exception e) {
            logger.info("AI调用异常: {}", e.getMessage(), e);

            if(e.getMessage().contains("401")) {
                chatModelFactory.remove(apiId);
            }
            throw new AIRelatedException(e.getMessage());
        }

        return Collections.singletonMap("generation", response);
    }

    /**
     * 不要问我什么原理，AI写的，作用是流式输出（就是像AI一个字一个字打出来的那样）
     * 没什么屌用，因为我最后没用这个方案，但是看起来很酷，所以就留着了
     * @return Flux
     */
    @PostMapping(value = "/generateStream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<StreamChunk> generateStream(@RequestBody Map<String, String> request, HttpSession session) {

        String apiId = request.get("apiId");
        String message = request.get("message");
        String apiKey = request.get("apiKey");
        String baseUrl = request.get("baseUrl");
        String model = request.get("model");
        double temperature = Integer.parseInt(request.get("temperature"));
        ChatModel chatModel = chatModelFactory.getOrCreate(apiId, apiKey, baseUrl, model, temperature);

        String sessionID = session.getId();
        conversationService.addUserMessage(sessionID, message);
        List<Message> history = conversationService.getOrCreateConversation(sessionID);

        return chatModel.stream(new Prompt(history))
                .map(response -> {
                    String deltaText = response.getResult().getOutput().getText();
                    boolean isComplete = "STOP".equals(response.getResult().getMetadata().getFinishReason());

                    if(isComplete){
                        conversationService.addAIResponse(sessionID, deltaText);
                    }
                    return new StreamChunk(
                            deltaText,
                            isComplete,
                            response.getResult().getMetadata().getFinishReason()
                    );
                })
                .filter(chunk -> !chunk.text().isEmpty());
    }

    @PostMapping("/remove")
    public String remove(@RequestBody String apiId){
        chatModelFactory.remove(apiId);
        return "SUCCESS";
    }
}
